import time
from typing import List

import ujson

from theHarvester.discovery.constants import MissingKey
from theHarvester.lib.core import AsyncFetcher, Core


class SearchPentestTools:
    def __init__(self, word) -> None:
        # Script is largely based off https://pentest-tools.com/public/api_client.py.txt
        self.word = word
        self.key = Core.pentest_tools_key()
        if self.key is None:
            raise MissingKey("PentestTools")
        self.total_results: List = []
        self.api = f"https://pentest-tools.com/api?key={self.key}"
        self.proxy = False

    async def poll(self, scan_id):
        while True:
            time.sleep(3)
            # Get the status of our scan
            scan_status_data = {"op": "get_scan_status", "scan_id": scan_id}
            responses = await AsyncFetcher.post_fetch(
                url=self.api, data=ujson.dumps(scan_status_data), proxy=self.proxy
            )
            res_json = ujson.loads(responses.strip())
            if res_json["op_status"] == "success":
                if (
                    res_json["scan_status"] != "waiting"
                    and res_json["scan_status"] != "running"
                ):
                    getoutput_data = {
                        "op": "get_output",
                        "scan_id": scan_id,
                        "output_format": "json",
                    }
                    responses = await AsyncFetcher.post_fetch(
                        url=self.api, data=ujson.dumps(getoutput_data), proxy=self.proxy
                    )

                    res_json = ujson.loads(responses.strip("\n"))
                    self.total_results = await self.parse_json(res_json)
                    break
            else:
                print(
                    f"Operation get_scan_status failed because: {res_json['error']}. {res_json['details']}"
                )
                break

    @staticmethod
    async def parse_json(json_results):
        status = json_results["op_status"]
        if status == "success":
            scan_tests = json_results["scan_output"]["output_json"]
            output_data = scan_tests[0]["output_data"]
            host_to_ip = [
                f"{subdomain[0]}:{subdomain[1]}"
                for subdomain in output_data
                if len(subdomain) > 0
            ]
            return host_to_ip
        return []

    async def get_hostnames(self) -> list:
        return self.total_results

    async def do_search(self) -> None:
        subdomain_payload = {
            "op": "start_scan",
            "tool_id": 20,
            "tool_params": {
                "target": f"{self.word}",
                "web_details": "off",
                "do_smart_search": "off",
            },
        }
        responses = await AsyncFetcher.post_fetch(
            url=self.api, data=ujson.dumps(subdomain_payload), proxy=self.proxy
        )
        res_json = ujson.loads(responses.strip())
        if res_json["op_status"] == "success":
            scan_id = res_json["scan_id"]
            await self.poll(scan_id)

    async def process(self, proxy: bool = False) -> None:
        self.proxy = proxy
        await self.do_search()  # Only need to do it once.
